
import copy
import torch
from thop import profile, clever_format
# from torchinfo import summary
from models.cnn import CNNCifar100
from models.MobileNetV2 import mobilenetv2
from models.ResNet_Dery import ResNet18_cifar10
from models.Vgg16 import VGG16

from algorithm.Training_FedDery4 import split_block

from models.model_creator import CombinedOneBlock

from sim.block_meta import MODEL_BLOCKS,MODEL_ZOO,PLANES
from utils.options import args_parser
args = args_parser()

net1 = CNNCifar100(args=args)
net2 = mobilenetv2(args=args)
net3 = ResNet18_cifar10(num_classes = args.num_classes,args=args)

# net1 = mobilenetv2(args)
# net2 = ResNet18_cifar10(num_classes=args.num_classes, args=args)
# net3 = VGG16(args)

net_zoo = []

net_zoo.append(net1)
net_zoo.append(net2)
net_zoo.append(net3)
num_model = len(net_zoo)

net_test = []
block_list = split_block(copy.deepcopy(net_zoo),4)
# for i in range(num_model):
#     net_test.append(CombinedOneBlock(block_list[i][0], block_list[i][1]).to(args.device))

net_half = CombinedOneBlock(block_list[2][2], block_list[2][3])

# 定义输入
input = torch.randn(50, 3, 32, 32)  # 输入大小为 (batch_size, channels, height, width)

# summary(net_half, input_size=(50, 3, 32, 32))

# net_glob = VGG16(args)
# 计算 FLOPs 和参数量
flops, params = profile(net_half, inputs=(input,))
flops, params = clever_format([flops, params], "%.3f")  # 格式化输出

print(f"Total FLOPs: {flops}")
print(f"Total parameters: {params}")


